package com.yanghui.dtracing.core.opentracing.propagation;

import cn.hutool.core.util.NumberUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.yanghui.dtracing.core.constant.AgentConstant;
import com.yanghui.dtracing.core.opentracing.DTracerSpanContext;
import io.opentracing.SpanContext;
import io.opentracing.propagation.TextMapExtract;
import io.opentracing.propagation.TextMapInject;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

public class TextMapJsonCodec implements Injector<TextMapInject>,Extractor<TextMapExtract>{
    @Override
    public DTracerSpanContext extract(TextMapExtract carrier) {
        Iterator<Map.Entry<String,String>> iterator = carrier.iterator();
        String contextJson = null;
        while(iterator.hasNext()){
            Map.Entry<String,String> entry = iterator.next();
            if(AgentConstant.HTTP_HEADER_SPAN_CONTEXT_KEY.equalsIgnoreCase(entry.getKey())){
                contextJson = entry.getValue();
                break;
            }
        }
        if(StrUtil.isEmpty(contextJson)){
            return null;
        }
        if(!JSONUtil.isJsonObj(contextJson)){
            return null;
        }
        JSONObject jsonObject = JSONUtil.parseObj(contextJson);
        String traceId = jsonObject.get("traceId",String.class);
        String spanId = jsonObject.get("spanId",String.class);
        if(StrUtil.isEmpty(traceId) || StrUtil.isEmpty(spanId)){
            return null;
        }
        Boolean sampled = jsonObject.get("sampled",Boolean.class);
        String parentId = jsonObject.get("parentId",String.class);;
        Map<String,String> baggages = jsonObject.get("baggageItems",Map.class);
        DTracerSpanContext context = new DTracerSpanContext(traceId,spanId,parentId,sampled);
        String childContextIndexStr = jsonObject.get("childContextIndex",String.class);
        if(!StrUtil.isEmpty(childContextIndexStr) && NumberUtil.isInteger(childContextIndexStr)){
            int childContextIndex = Integer.valueOf(childContextIndexStr);
            if(childContextIndex > 0){
                context.setChildContextIndex(new AtomicInteger(childContextIndex - 1));
            }
        }
        context.addBaggageItems(baggages);
        return context;
    }

    @Override
    public void inject(SpanContext spanContext, TextMapInject carrier) {
        DTracerSpanContext context = (DTracerSpanContext)spanContext;
        Map<String,Object> map = new HashMap<>();
        map.put("traceId",context.toTraceId());
        map.put("spanId",context.toSpanId());
        map.put("sampled",String.valueOf(context.isSampled()));
        map.put("parentId",context.getParentId());
        map.put("childContextIndex",context.getChildContextIndex().incrementAndGet());
        Map<String,String> baggage = ((DTracerSpanContext) spanContext).getBaggageItems();
        map.put("baggageItems",baggage);
        carrier.put(AgentConstant.HTTP_HEADER_SPAN_CONTEXT_KEY, JSONUtil.toJsonStr(map));
    }
}
